from typing import Sequence, Collection, Union

import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import accuracy_score

import numpy as np
from PIL import Image
from tqdm import tqdm

from modules import DigitalLetterNet
from dataset import DIGITS_LETTERS
from utils import ver_code_split


class VerCodeIdentifier(object):
    def __init__(self, in_chs=3, out_chs=36, pretrained: str = None, fcn=False, device='cuda'):
        self.in_chs = in_chs
        self.out_chs = out_chs
        self.device = torch.device(device)

        self.net: nn.Module = DigitalLetterNet(self.in_chs, self.out_chs, fcn).to(self.device)
        if pretrained:
            self.net.load_state_dict(torch.load(pretrained, map_location=self.device))

    def fit(self, dataset: Dataset, epochs=10, lr=0.001, batch_size=512, num_workers=4):
        criterion = nn.CrossEntropyLoss()
        optimizer = Adam(self.net.parameters(), lr=lr)
        self.net.train()
        for epoch in range(epochs):
            acc_lst = []
            loss_lst = []
            train_loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
            data_iter = tqdm(train_loader, unit='batch', postfix={'epoch': 0, 'loss': 0, 'acc': 0})
            for batch, label in data_iter:
                optimizer.zero_grad()
                batch = batch.permute((0, 3, 1, 2)).float().to(self.device)
                label = label.long().to(self.device)
                output = self.net(batch)
                loss = criterion(output, label)

                loss.backward()  # 反向传播，计算梯度
                optimizer.step()  # 更新网络参数

                # 计算准确率
                out_label = torch.argmax(output.detach(), dim=1).cpu()
                acc = accuracy_score(out_label.numpy(), label.cpu().numpy())
                acc_lst.append(acc)
                loss_lst.append(loss.item())
                data_iter.set_postfix(
                    epoch='{:0>2d}'.format(epoch + 1),
                    loss="{:.6f}({:.6f})".format(loss.item(), np.mean(loss_lst)),
                    acc="{:.6f}({:.6f})".format(acc, np.mean(acc_lst))
                )

    @torch.no_grad()
    def predict(self, data: Union[np.ndarray, torch.Tensor, Sequence, Collection]):
        self.net.eval()
        if isinstance(data, (np.ndarray, torch.Tensor)):
            data = torch.tensor(data, dtype=torch.float)
        elif isinstance(data, (Sequence, Collection)):
            data = torch.stack([torch.tensor(d, dtype=torch.float) for d in data], 0)
        else:
            raise TypeError('Unspported type:%s' % (type(data)))

        assert 3 <= data.dim() <= 4, 'data dim must be 3 or 4'
        if data.dim() == 3:
            data.unsqueeze_(0)
        output = self.net(data.to(self.device)).detach().cpu()
        return output.argmax(dim=-1).squeeze().numpy()

    def identify(self, img: Union[str, Image]) -> str:
        """识别验证码
        :param img: str/Image, 验证码路径或图像
        :return: str, 验证码字符
        """
        ver_imgs = ver_code_split(img)
        ver_pred = self.net.predict(ver_imgs)
        label_pred = ''.join([DIGITS_LETTERS[vp] for vp in ver_pred])
        return label_pred
